from matplotlib.pylab import date2num
import baostock as bs
import pandas as pd
import datetime
import ZIndex.stock.index as index

stock_list = ['sz.300008', 'sz.002623', 'sz.000665', 'sz.002411', 'sz.300222']

def date_to_num(dates):
    num_time = []
    for date in dates:
        date_time = datetime.datetime.strptime(date, '%Y-%m-%d')
        num_date = date2num(date_time)
        num_time.append(num_date)
    return num_time

# 绘制蜡烛图
def format_date(x, pos=None):
    # 日期格式化函数，根据天数索引取出日期值
    return '' if x < 0 or x > len(date_tickers) - 1 else date_tickers[int(x)]


lg = bs.login()
for i in stock_list:
    print(i)
    rs = bs.query_history_k_data_plus(i,
                                      "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                      start_date='2019-01-01', end_date='2021-12-31',
                                      frequency="d", adjustflag="3")

    #### 打印结果集 ####
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    # 二维数组
    result = result.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume','amount']]

    date_tickers = result.date.values
    result.date = range(0, len(result))  # 日期改变成序号
    matix = result.values  # 转换成绘制蜡烛图需要的数据格式(date, open, close, high, low, volume)
    xdates = matix[:, 0]  # X轴数据(这里用的天数索引)
    # print(result['volume'])

    # adReal = talib.AD(result['high'], result['low'], result['close'], result['volume'])
    # ma5Price = talib.MA(result['close'], timeperiod=5, matype=0)
    # sarPrice = talib.SAR(result['high'], result['low'], acceleration=0, maximum=0)
    """ 参数设置 begin """

    """ 参数设置 end """
    # KDJ指标
    result['kdj_w'], result['kdj_d'], result['kdj_j'] = index.KDJIndex(result)
    # WR指标
    result['wr'] = index.WR(result)
    # MACD 指标
    result['macd'] = index.MACD(result['close'])
    # ARBR 指标
    result['ar'], result['br'] = index.arbr(result)
    # CYR 指标
    result['CYR'] = index.CYR(result)

    print("--result")
    print(result[-1:].T)






